import os
os.environ['STREAMLIT_BROWSER_GATHER_USAGE_STATS'] = "false"

import streamlit as st
import pandas as pd
from typing import Dict, List
from src.core.config import get_assistant_config
from src.core.test_checker import check_test_cases
from src.core.file_processor import process_file
from src.core.test_point import generate_test_points
from src.core.test_case import generate_test_cases
from src.core.excel import generate_excel
from src.utils.image import get_base64_image

# 设置页面配置
st.set_page_config(
    page_title="AI 测试用例生成",
    page_icon="../assets/header-logo.png",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义CSS样式
st.markdown("""
<style>
    .main {
        max-width: 1200px;
        padding: 2rem;
    }
    /* 让进度条占满宽度 */
    .element-container, .stMarkdown, .stText {
        width: 100% !important;
        max-width: 100% !important;
    }
    /* 进度条容器样式 */
    .stProgress > div {
        width: 100% !important;
        max-width: 100% !important;
    }
    /* 代码块和JSON显示样式 */
    .stCodeBlock, pre {
        width: 100% !important;
        max-width: 100% !important;
    }
    /* 文本输出样式 */
    .fixed-width {
        font-family: monospace;
        white-space: pre-wrap;
        width: 100% !important;
        max-width: 100% !important;
    }
    /* 按钮样式 */
    .stButton>button {
        width: 100%;
        margin-top: 1rem;
    }
    .stProgress .st-bo {
        background-color: #00a0dc;
    }
    .stDataFrame {
        margin-top: 1rem;
    }
    .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
        max-width: 100% !important;
    }
    h1 {
        color: #2c3e50;
        margin-bottom: 2rem;
    }
    h2 {
        font-size: 1.5rem;
        color: #34495e;
        margin-top: 2rem;
    }
    .stAlert {
        margin-top: 1rem;
        width: 100% !important;
    }
    /* Expander 样式 */
    .streamlit-expanderHeader {
        width: 100% !important;
        max-width: 100% !important;
    }
    .streamlit-expanderContent {
        width: 100% !important;
        max-width: 100% !important;
        padding: 0 !important;
    }
    /* 添加标题图标样式 */
    .title-container {
        display: flex;
        flex-direction: column;  /* 整体垂直排列 */
        align-items: center;     /* 水平居中 */
        gap: 5px;               /* 行间距 */
        padding: 1rem 0;
        margin-bottom: 1.5rem;
    }
    
    /* 添加标题行容器 */
    .title-row {
        display: flex;          /* 水平排列 */
        align-items: center;    /* 垂直居中对齐 */
        gap: 10px;             /* 图标和标题间距 */
    }

    .title-icon {
        width: 40px;          /* 固定宽度 */
        height: auto;          /* 高度自动调整 */
        object-fit: contain;   /* 保持长宽比 */
    }

    .title-container h1 {
        font-size: 2rem;       /* 与首页保持一致 */
        color: #2c3e50;
        margin: 0;
    }
</style>
""", unsafe_allow_html=True)

def init_session_state():
    """初始化session state"""
    if "results" not in st.session_state:
        st.session_state.results = None
    if "test_points" not in st.session_state:
        st.session_state.test_points = None
    if "excel_data" not in st.session_state:
        st.session_state.excel_data = None

def process_requirement(file, selected_models: List[str], progress_bar, status_text) -> Dict:
    """处理需求文档并生成测试用例"""
    try:
        with st.expander("处理详情", expanded=True):
            # 1. 解析文件 (20%)
            status_text.text("1/5 解析需求文档...")
            st.markdown("### 1. 解析需求文档")
            content = process_file(file)
            st.markdown("**需求文档内容：**")
            st.code(content, language='markdown')
            progress_bar.progress(20)
            
            # 2. 分析测试点 (40%)
            status_text.text("2/5 分析测试点...")
            st.markdown("### 2. 分析测试点")
            analyzer_config = get_assistant_config('requirement_analyzer')
            st.markdown(f"**使用需求分析助手:** {analyzer_config['name']}")
            #st.markdown(f"**{selected_models}**")
            test_points = generate_test_points(content, analyzer_config['api_key'], selected_models)
            st.markdown("**生成的测试点：**")
            st.json(test_points)
            progress_bar.progress(40)
            
            # 3. 生成测试用例 (60%)
            status_text.text("3/5 生成测试用例...")
            st.markdown("### 3. 生成测试用例")
            test_generator_config = get_assistant_config('test_generator')
            st.markdown(f"**使用测试用例生成助手:** {test_generator_config['name']}")
            test_cases, category_stats = generate_test_cases(test_points, test_generator_config['api_key'], selected_models)
            if not test_cases:
                raise Exception("未能生成任何测试用例")
            st.markdown("**生成的测试用例：**")
            st.json(test_cases)
            progress_bar.progress(60)
            
            # 4. 检查用例质量 (80%)
            status_text.text("4/5 检查用例质量...")
            st.markdown("### 4. 检查用例质量")
            checker_config = get_assistant_config('test_checker')
            st.markdown(f"**使用测试用例检查助手:** {checker_config['name']}")
            test_cases = check_test_cases(test_cases, checker_config['api_key'])
            progress_bar.progress(80)
            
            # 5. 生成Excel (100%)
            status_text.text("5/5 优化并导出...")
            st.markdown("### 5. 生成Excel")
            excel_data = generate_excel(test_cases)
            if not excel_data:
                raise Exception("Excel生成失败")
            st.success("Excel文件生成成功！")
            progress_bar.progress(100)
            
            return {
                "test_points": test_points,
                "test_cases": test_cases,
                "excel_data": excel_data,
                "category_stats": category_stats
            }
            
    except Exception as e:
        st.error(f"处理过程中出现错误: {str(e)}")
        return None

def create_preview_df(test_cases: List[Dict]) -> pd.DataFrame:
    """创建预览数据框"""
    df = pd.DataFrame(test_cases)
    
    # 定义要显示的列及其顺序
    preview_columns = [
        '用例名称', 
        '所属模块',
        '用例等级',
        '前置条件',
        '测试步骤',
        '预期结果',
        '标签'
    ]
    
    # 确保所有列都存在
    for col in preview_columns:
        if col not in df.columns:
            df[col] = ''
            
    # 处理列表类型的字段
    list_columns = ['前置条件', '测试步骤', '预期结果', '标签']
    for col in list_columns:
        if col in df.columns:
            df[col] = df[col].apply(lambda x: '\n'.join(x) if isinstance(x, list) else x)
    
    return df[preview_columns]

def main():
    # 使用 HTML 来组合图标和标题
    st.markdown(
        """
        <div class="title-container">
            <div class="title-row">
                <img src="data:image/png;base64,{}" class="title-icon">
                <h1>测试用例生成</h1>
            </div>
        </div>
        """.format(get_base64_image("../assets/header-logo.png")),
        unsafe_allow_html=True
    )
    # 初始化session state
    init_session_state()
    
    # 1. 配置区
    with st.sidebar:
        st.header("⚙️ 配置")
        
        # 测试类别选择
        st.subheader("测试用例类别")
        selected_models = st.multiselect(
            "选择要生成的测试用例类别",
            ["功能测试", "性能测试", "安全测试", "兼容性测试", "可用性测试", "安装测试", "配置测试", "探索性测试", "自动化测试", "灾难恢复测试", "接口测试", "本地化测试", "负载测试", "压力测试", "容量测试"],
            default=["功能测试", "性能测试", "安全测试", "接口测试","压力测试", "可用性测试"],
            help="可以选择多个测试用例类别"
        )
        
        # 显示配置信息
        if selected_models:
            st.info(f"已选择 {len(selected_models)} 种测试用例类别")
    
    # 2. 文件上传区
    #st.header("📄 需求文档")
    st.markdown("<h2 style='font-size: 24px;'>📄 需求文档</h2>", unsafe_allow_html=True)
    uploaded_file = st.file_uploader(
        "上传需求文档",
        type=["md", "docx", "txt"],
        help="支持 Markdown、Word 和文本文件"
    )
    
    # 3. 处理区
    if uploaded_file:
        if st.button("🚀 生成测试用例", use_container_width=True):
            try:
                with st.spinner("🔄 正在处理..."):
                    # 显示进度
                    progress_container = st.container()
                    with progress_container:
                        progress_bar = st.progress(0)
                        status_text = st.empty()
                        
                        # 处理文件并更新进度
                        results = process_requirement(uploaded_file, selected_models, progress_bar, status_text)
                        
                        if results:
                            st.session_state.results = results
                            st.session_state.test_points = results["test_points"]
                            st.session_state.excel_data = results["excel_data"]
                            status_text.text("✅ 测试用例生成完成！")
                            
            except Exception as e:
                st.error(f"❌ 处理失败: {str(e)}")
    
    # 4. 结果预览区
    if st.session_state.get("results"):
        st.header("📊 生成结果")
        
        # 测试点预览
        if st.session_state.test_points:
            with st.expander("🎯 测试点预览"):
                st.json(st.session_state.test_points)
        
        # Excel预览
        if st.session_state.results.get("test_cases"):
            st.subheader("📋 测试用例预览")
            
            # 创建预览表格
            preview_df = create_preview_df(st.session_state.results["test_cases"])
            
            # 显示统计信息
            col1, col2 = st.columns(2)
            with col1:
                st.metric("总用例数", len(preview_df))
            with col2:
                category_count = len(selected_models)
                st.metric("覆盖测试类别数", category_count)
            
            # 显示表格
            st.dataframe(
                preview_df,
                height=600,
                use_container_width=True
            )
            
            # 按测试类别统计
            with st.expander("📊 测试类别统计"):
                category_stats = st.session_state.results["category_stats"]
                
                # 格式化显示
                st.markdown("### 测试用例分布")
                total_cases = len(st.session_state.results["test_cases"])
                
                # 过滤选中的类别
                filtered_stats = {k: v for k, v in category_stats.items() if k in selected_models}
                
                if filtered_stats:  # 只在有数据时显示
                    # 创建等分的列
                    cols = st.columns(len(filtered_stats))
                    
                    # 在每列中显示对应类别的统计
                    for col, (category, count) in zip(cols, filtered_stats.items()):
                        with col:
                            st.metric(
                                label=category,
                                value=f"{count} 个用例",
                                delta=f"占比 {count/total_cases*100:.1f}%",
                                help=f"该类别包含 {count} 个测试用例，占总数的 {count/total_cases*100:.1f}%"
                            )
                    
                    # 显示柱状图
                    if len(filtered_stats) > 0:
                        chart_data = pd.Series(filtered_stats)
                        if not chart_data.empty:
                            st.bar_chart(chart_data)
                else:
                    st.info("所选测试类别暂无数据")
        
        # 下载按钮
        if st.session_state.excel_data:
            st.download_button(
                "⬇️ 下载Excel文件",
                data=st.session_state.excel_data,
                file_name="test_cases.xlsx",
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                use_container_width=True
            )

if __name__ == "__main__":
    main() 